#ifndef __MODULE__
#define __MODULE__

#include <vector>
#include <memory>
#include <iostream>
#include "../../Tensor.hpp"



namespace DimN
{

    template<typename _T = float>
    class Module {
    public:
        Module() {}

        // Virtual destructor to ensure proper cleanup of derived classes
        virtual ~Module() {}

        // Pure virtual forward function that must be overridden by all subclasses
        virtual DimN::Tensor<_T> forward(DimN::Tensor<_T> x) {
            return DimN::Tensor<_T>();
        };

        // Add a method to register parameters if needed
        void register_parameter(const std::string& name, std::shared_ptr<double> parameter) {
            parameters.push_back(parameter);
            parameter_names.push_back(name);
        }

        // Optionally, add methods for handling submodules if your design requires
        void add_module(const std::string& name, std::shared_ptr<Module> module) {
            modules.push_back(module);
            module_names.push_back(name);
        }

    private:
        std::vector<std::shared_ptr<double>> parameters;
        std::vector<std::string> parameter_names;
        std::vector<std::shared_ptr<Module>> modules;
        std::vector<std::string> module_names;

        // Base class for all neural network modules.
        // Your models should also subclass this class.
        // Modules can also contain other Modules, allowing to nest them in a tree structure.
    };

};

#endif // __MODULE__